#include "DatabaseManager.h"
#include "DatabaseManagerException.h"

#include <QDate>
#include <QStringList>
#include <QSqlQuery>
#include <QSqlError>
#include <QDebug>
#include <QVariant>
#include <assert.h>

using namespace kryptos::entities;

namespace kryptos { namespace database {

DatabaseManager::DatabaseManager()
{
}

DatabaseManager::~DatabaseManager()
{
    m_keyRingDB.close();
    m_groupsDB.close();
}

DatabaseManager& DatabaseManager::manager()
{
    static DatabaseManager instance;
    return instance;
}

void DatabaseManager::openKeyring()
{
    if (!m_keyRingDB.isOpen())
    {
        m_keyRingDB = QSqlDatabase::addDatabase("QSQLITE", "keyring");
        m_keyRingDB.setDatabaseName("kryptos_keyring.k");

        if (!m_keyRingDB.open())
            //qDebug() << "SQL ERROR OPENING: '" << m_keyRingDB.lastError().text() << "'";
            throw DatabaseManagerException(DatabaseManagerException::OPEN_ERROR, m_keyRingDB.lastError().text());

        QSqlQuery query(m_keyRingDB);
        QString str = "CREATE TABLE IF NOT EXISTS keyring("
                      "keyid       TEXT NOT NULL,"
                      "name        TEXT NOT NULL,"
                      "email       TEXT NOT NULL,"
                      "comment     TEXT NOT NULL,"
                      "type        INT NOT NULL,"
                      "size        INT NOT NULL,"
                      "hashphrase  TEXT NOT NULL,"
                      "timestamp   TEXT NOT NULL,"
                      "publickey   BLOB,"
                      "privatekey  BLOB,"
                      "PRIMARY KEY(keyid));";

        if (!query.exec(str))
            throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, query.lastError().text() + ". Query: '" + query.lastQuery() + "'");
    }
}

void DatabaseManager::openGroups()
{
    if (!m_groupsDB.isOpen())
    {
        m_groupsDB = QSqlDatabase::addDatabase("QSQLITE", "usergroups");
        m_groupsDB.setDatabaseName("kryptos_groups.k");

        if (!m_groupsDB.open())
            qDebug() << "SQL ERROR OPENING: '" << m_groupsDB.lastError().text() << "'";

        QSqlQuery query(m_groupsDB);
        QString str = "CREATE TABLE IF NOT EXISTS groupmembers("
                      "groupname   TEXT NOT NULL,"
                      "keyid       TEXT NOT NULL,"
                      "symkeyid    TEXT ,"
                      "PRIMARY KEY(groupname, keyid, symkeyid));";

        if (!query.exec(str))
            throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, query.lastError().text() + ". Query: '" + query.lastQuery() + "'");

        str =         "CREATE TABLE IF NOT EXISTS keyfiles("
                      "symkeyid       TEXT NOT NULL,"
                      "symkeyfilename TEXT NOT NULL,"
                      "decryptonly    BOOLEAN,"
                      "offset         UNSIGNED BIG INT,"
                      "expiration     DATE,"
                      "PRIMARY KEY(symkeyid));";

        if (!query.exec(str))
            throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, query.lastError().text() + ". Query: '" + query.lastQuery() + "'");

        str =         "CREATE TABLE IF NOT EXISTS groups("
                      "groupname      TEXT  NOT NULL,"
                      "type           INT   NOT NULL,"
                      "PRIMARY KEY(groupname));";

        if (!query.exec(str))
            throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, query.lastError().text() + ". Query: '" + query.lastQuery() + "'");
    }
}

void DatabaseManager::addKeyPair(const KeyInfo& info, const QByteArray& privatePart, const QByteArray& publicPart)
{
    QSqlQuery query(m_keyRingDB);

    query.prepare("INSERT INTO keyring (keyid, name, email, comment, type, size, hashphrase, timestamp, publickey, privatekey) "
                  "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)");

    query.addBindValue(info.keyid().c_str());
    query.addBindValue(info.name().c_str());
    query.addBindValue(info.email().c_str());
    query.addBindValue(info.comment().c_str());
    query.addBindValue(info.type());
    query.addBindValue(info.size());
    query.addBindValue(info.passphrase().c_str());
    query.addBindValue(QDate::currentDate());
    query.addBindValue(publicPart, QSql::In | QSql::Binary);
    query.addBindValue(privatePart, QSql::In | QSql::Binary);

    if (!query.exec())
        //qDebug() << "SQL ERROR EXECUTING QUERY: '" << m_keyRingDB.lastError().text() << "'";
        throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, m_keyRingDB.lastError().text() + ". Query: '" + query.lastQuery() + "'");
}

void DatabaseManager::addPublicKey(const KeyInfo& info, const QByteArray& publicPart)
{
    QSqlQuery query(m_keyRingDB);

    query.prepare("INSERT INTO keyring (keyid, name, email, comment, type, size, hashphrase, timestamp, publickey, privatekey) "
                  "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)");

    query.addBindValue(info.keyid().c_str());
    query.addBindValue(info.name().c_str());
    query.addBindValue(info.email().c_str());
    query.addBindValue(info.comment().c_str());
    query.addBindValue(info.type());
    query.addBindValue(info.size());
    query.addBindValue(info.passphrase().c_str());
    query.addBindValue(QDate::currentDate());
    query.addBindValue(publicPart, QSql::In | QSql::Binary);
    query.addBindValue(QByteArray(), QSql::In | QSql::Binary);

    if (!query.exec())
        throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, m_keyRingDB.lastError().text() + ". Query: '" + query.lastQuery() + "'");
}

KeyInfoList DatabaseManager::allKeys()
{
    QScopedPointer<QSqlQuery> result(searchKey());

    KeyInfoList infolist;

    while (result->next())
    {
        QString keyid         = result->value(0).toString();
        QString name          = result->value(1).toString();
        QString email         = result->value(2).toString();
        QString comment       = result->value(3).toString();
        KeyInfo::KeyType type = (KeyInfo::KeyType)result->value(4).toInt();
        int size              = result->value(5).toInt();
        QString hashphrase    = result->value(6).toString();
        QDate timestamp       = result->value(7).toDate();
        QByteArray publicKey  = result->value(8).toByteArray();
        QByteArray privateKey = result->value(9).toByteArray();

        KeyInfo info(type,
                     keyid.toStdString(),
                     name.toStdString(),
                     email.toStdString(),
                     comment.toStdString(),
                     hashphrase.toStdString(),
                     timestamp.toString().toStdString(),
                     size);

        if (!publicKey.isEmpty() && !privateKey.isEmpty())
            info.setIsIdentity(true);

        infolist.push_back(info);

    }

    return infolist;
}

KeyInfoList DatabaseManager::keysWithPublicAndPrivateData()
{
    QScopedPointer<QSqlQuery> result(searchKey());

    KeyInfoList infolist;

    while (result->next())
    {
        // Only results that have public and private keys
        if (!result->value(9).isNull())
        {
            QString keyid         = result->value(0).toString();
            QString name          = result->value(1).toString();
            QString email         = result->value(2).toString();
            QString comment       = result->value(3).toString();
            KeyInfo::KeyType type = (KeyInfo::KeyType)result->value(4).toInt();
            int size              = result->value(5).toInt();
            QString hashphrase    = result->value(6).toString();
            QDate timestamp       = result->value(7).toDate();

            KeyInfo info(type,
                         keyid.toStdString(),
                         name.toStdString(),
                         email.toStdString(),
                         comment.toStdString(),
                         hashphrase.toStdString(),
                         timestamp.toString().toStdString(),
                         size);

            info.setIsIdentity(true);

            infolist.push_back(info);
        }
    }

    return infolist;
}

QSqlQuery* DatabaseManager::searchKey(QMap<QString, QString> params)
{
    QSqlQuery* query = new QSqlQuery(m_keyRingDB);

    // builds WHERE condition line with params from map
    // if map is empty, query will search for all keys
    QString conditions;
    QMapIterator<QString, QString> it(params);
    while (it.hasNext())
    {
        it.next();
        if (conditions.isEmpty())
            conditions += " WHERE ";
        else
            conditions += " AND ";

        conditions += it.key()+ " = '" + it.value() + "'";
    }

    if (!query->exec("SELECT keyid, name, email, comment, type, size, hashphrase, timestamp, publickey, privatekey "
                     "FROM keyring " +
                     conditions + ";"))
        throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, m_keyRingDB.lastError().text() + ". Query: '" + query->lastQuery() + "'");

    return query;
}

bool DatabaseManager::hasKeyWithEmail(const QString& email)
{
    QMap<QString, QString> params;
    params["email"] = email;

    QScopedPointer<QSqlQuery> result(searchKey(params));

    return result->first();
}

bool DatabaseManager::hasKeyWithId(const QString& id)
{
    QMap<QString, QString> params;
    params["keyid"] = id;

    QScopedPointer<QSqlQuery> result(searchKey(params));

    return result->first();
}

bool DatabaseManager::hasKeyWithNameAndEmail(const QString& name, const QString& email)
{
    QMap<QString, QString> params;
    params["name"] = name;
    params["email"] = email;

    QScopedPointer<QSqlQuery> result(searchKey(params));

    return result->first();
}

KeyDAO DatabaseManager::getKey(const QString& idOrEmail)
{
    KeyDAO dao;
    QString searchField = idOrEmail.contains('@') ? "email" : "keyid";

    QMap<QString, QString> params;
    params[searchField] = idOrEmail;

    QScopedPointer<QSqlQuery> result(searchKey(params));

    if (result->first())
    {
        QString keyid         = result->value(0).toString();
        QString name          = result->value(1).toString();
        QString email         = result->value(2).toString();
        QString comment       = result->value(3).toString();
        KeyInfo::KeyType type = (KeyInfo::KeyType)result->value(4).toInt();
        int size              = result->value(5).toInt();
        QString hashphrase    = result->value(6).toString();
        QDate timestamp       = result->value(7).toDate();
        QByteArray publicKey  = result->value(8).toByteArray();
        QByteArray privateKey = result->value(9).toByteArray();

        KeyInfo info(type,
                     keyid.toStdString(),
                     name.toStdString(),
                     email.toStdString(),
                     comment.toStdString(),
                     hashphrase.toStdString(),
                     timestamp.toString().toStdString(),
                     size);

        if (!publicKey.isEmpty() && !privateKey.isEmpty())
            info.setIsIdentity(true);

        dao = KeyDAO(info, privateKey, publicKey);
    }

    return dao;
}

//////////////////////////////////////////////////////////////////////////////////
///////////////////////////////// CONTACTS AND GROUPS ////////////////////////////
//////////////////////////////////////////////////////////////////////////////////

void DatabaseManager::addGroup(const Group& group)
{
    QSqlQuery query(m_groupsDB);

    query.prepare("INSERT INTO groups (groupname, type) "
                  "VALUES (?, ?)");

    query.bindValue(0, group.name().c_str());
    query.bindValue(1, group.type());

    if (!query.exec())
        throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, query.lastError().text() + ". Query: '" + query.lastQuery() + "'");

    foreach (const Contact& member, group.members())
    {
        foreach (const SymmetricKeyInfo& symkeyinfo, member.symmetricKeysInfo())
        {
            query.prepare("INSERT INTO groupmembers (groupname, keyid, symkeyid) "
                          "VALUES (?, ?, ?)");

            query.bindValue(0, group.name().c_str());
            query.bindValue(1, member.keyInfo().keyid().c_str());
            query.bindValue(2, symkeyinfo.id().c_str());

            if (!query.exec())
                throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, query.lastError().text() + ". Query: '" + query.lastQuery() + "'");


            query.prepare("INSERT INTO keyfiles (symkeyid, symkeyfilename, decryptonly, offset, expiration) "
                          "VALUES (?, ?, ?, ?, ?)");

            query.bindValue(0, symkeyinfo.id().c_str());
            query.bindValue(1, symkeyinfo.keyfile().c_str());
            query.bindValue(2, symkeyinfo.isDecryptOnly());
            query.bindValue(3, symkeyinfo.offset());
            query.bindValue(4, 0);

            if (!query.exec())
                throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, query.lastError().text() + ". Query: '" + query.lastQuery() + "'");
        }
    }
}

bool DatabaseManager::hasGroup(const std::string& name)
{
    QSqlQuery query(m_groupsDB);

    if (!query.exec("SELECT groupname, type FROM groups WHERE groupname = '" + QString(name.c_str()) + "';"))
        throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, m_groupsDB.lastError().text() + ". Query: '" + query.lastQuery() + "'");

    return query.first();
}

void DatabaseManager::addContactToGroup(const Contact& contact, const Group& group)
{
    QSqlQuery query(m_groupsDB);


    if (!contact.hasSymmetricKey())
    {
        query.prepare("INSERT INTO groupmembers (groupname, keyid, symkeyid) "
                      "VALUES (?, ?, 0)");

        query.bindValue(0, group.name().c_str());
        query.bindValue(1, contact.keyInfo().keyid().c_str());

        if (!query.exec())
            throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, query.lastError().text() + ". Query: '" + query.lastQuery() + "'");
    }
    else
    {
        foreach (const SymmetricKeyInfo& symkeyinfo, contact.symmetricKeysInfo())
        {
            query.prepare("INSERT INTO groupmembers (groupname, keyid, symkeyid) "
                          "VALUES (?, ?, ?)");

            query.bindValue(0, group.name().c_str());
            query.bindValue(1, contact.keyInfo().keyid().c_str());
            query.bindValue(2, symkeyinfo.id().c_str());

            if (!query.exec())
                throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, query.lastError().text() + ". Query: '" + query.lastQuery() + "'");


            query.prepare("INSERT INTO keyfiles (symkeyid, symkeyfilename, decryptonly, offset, expiration) "
                          "VALUES (?, ?, ?, ?, ?)");

            query.bindValue(0, symkeyinfo.id().c_str());
            query.bindValue(1, symkeyinfo.keyfile().c_str());
            query.bindValue(2, symkeyinfo.isDecryptOnly());
            query.bindValue(3, symkeyinfo.offset());
            query.bindValue(4, 0);

            if (!query.exec())
                throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, query.lastError().text() + ". Query: '" + query.lastQuery() + "'");
        }
    }
}

Group DatabaseManager::findGroup(const std::string& name)
{
    QSqlQuery query(m_groupsDB);

    if (!query.exec("SELECT groups.groupname, groups.type, "
                            "groupmembers.groupname, groupmembers.keyid, groupmembers.symkeyid, "
                            "keyfiles.symkeyid, keyfiles.symkeyfilename, keyfiles.decryptonly, keyfiles.offset "
                    "FROM   groups, groupmembers, keyfiles "
                    "WHERE  groups.groupname = groupmembers.groupname "
                    "AND    groups.groupname = '" + QString(name.c_str()) + "' "
                    "AND    groupmembers.symkeyid = keyfiles.symkeyid;"))
        throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, m_groupsDB.lastError().text() + ". Query: '" + query.lastQuery() + "'");

    QMap<QString, KeyInfo> infomap;
    KeyInfoList keys = allKeys();
    foreach(const KeyInfo& info, keys)
    {
        infomap[info.keyid().c_str()] = info;
    }

    QMap<QString, Contact> members;
    Group::GroupType type;

    while (query.next())
    {
        QString groupname         = query.value(0).toString();
        type                      = (Group::GroupType)query.value(1).toInt();
        QString keyid             = query.value(3).toString();
        QString symkeyid          = query.value(4).toString();
        QString symkeyfile        = query.value(6).toString();
        bool decryptOnly          = query.value(7).toBool();
        unsigned long long offset = query.value(8).toULongLong();

        if (!members.contains(keyid))
        {
            Contact c(infomap.value(keyid));
            c.addSymmetricKey(SymmetricKeyInfo(symkeyid.toStdString(), symkeyfile.toStdString(), decryptOnly, offset));
            members[keyid] = c;
        }
        else
        {
            members[keyid].addSymmetricKey(SymmetricKeyInfo(symkeyid.toStdString(), symkeyfile.toStdString(), decryptOnly, offset));
        }
    }

    Group group(name, type);
    QMapIterator<QString, Contact> it(members);
    while (it.hasNext())
    {
        it.next();
        group.addMember(it.value());
    }

    return group;
}

GroupList DatabaseManager::allGroups()
{
    GroupList list  = allGroupsWithoutKeyFiles();
    GroupList list2 = allGroupsWithKeyFiles();

    std::copy(list2.begin(), list2.end(), std::inserter(list, list.end()));

    return list;
}

GroupList DatabaseManager::allGroupsWithoutKeyFiles()
{
    QSqlQuery query(m_groupsDB);

    if (!query.exec("SELECT groups.groupname, groups.type, "
                            "groupmembers.groupname, groupmembers.keyid, groupmembers.symkeyid "
                    "FROM   groups, groupmembers "
                    "WHERE  groups.groupname = groupmembers.groupname "
                    "AND    groupmembers.symkeyid = 0;"))
        throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, m_groupsDB.lastError().text() + ". Query: '" + query.lastQuery() + "'");

    QMap<QString, KeyInfo> infomap;
    KeyInfoList keys = allKeys();
    foreach(const KeyInfo& info, keys)
    {
        infomap[info.keyid().c_str()] = info;
    }

    QMap<QString, Group> groups;

    while (query.next())
    {
        QString groupname         = query.value(0).toString();
        Group::GroupType type     = (Group::GroupType)query.value(1).toInt();
        QString keyid             = query.value(3).toString();

        Contact c(infomap.value(keyid));

        if (!groups.contains(groupname))
        {
            Group g(groupname.toStdString(), type);
            g.addMember(c);
            groups[groupname] = g;
        }
        else
        {
            groups[groupname].addMember(c);
        }
    }

    return groups.values().toVector().toStdVector();
}

GroupList DatabaseManager::allGroupsWithKeyFiles()
{
    QSqlQuery query(m_groupsDB);

    if (!query.exec("SELECT groups.groupname, groups.type, "
                            "groupmembers.groupname, groupmembers.keyid, groupmembers.symkeyid, "
                            "keyfiles.symkeyid, keyfiles.symkeyfilename, keyfiles.decryptonly, keyfiles.offset "
                    "FROM   groups, groupmembers, keyfiles "
                    "WHERE  groups.groupname = groupmembers.groupname "
                    "AND    groupmembers.symkeyid = keyfiles.symkeyid;"))
        throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, m_groupsDB.lastError().text() + ". Query: '" + query.lastQuery() + "'");

    QMap<QString, KeyInfo> infomap;
    KeyInfoList keys = allKeys();
    foreach(const KeyInfo& info, keys)
    {
        infomap[info.keyid().c_str()] = info;
    }

    QMap<QString, Group> groups;
    QMap<QString, Contact> members;
    QMap<QString, QString> memberToGroupMap;

    while (query.next())
    {
        QString groupname         = query.value(0).toString();
        Group::GroupType type     = (Group::GroupType)query.value(1).toInt();
        QString keyid             = query.value(3).toString();
        QString symkeyid          = query.value(4).toString();
        QString symkeyfile        = query.value(6).toString();
        bool decryptOnly          = query.value(7).toBool();
        unsigned long long offset = query.value(8).toULongLong();

        if (!members.contains(keyid))
        {
            Contact c(infomap.value(keyid));
            c.addSymmetricKey(SymmetricKeyInfo(symkeyid.toStdString(), symkeyfile.toStdString(), decryptOnly, offset));
            members[keyid] = c;
            memberToGroupMap[keyid] = groupname;
        }
        else
        {
            members[keyid].addSymmetricKey(SymmetricKeyInfo(symkeyid.toStdString(), symkeyfile.toStdString(), decryptOnly, offset));
        }

        if (!groups.contains(groupname))
        {
            Group g(groupname.toStdString(), type);
            groups[groupname] = g;
        }
    }

    QMapIterator<QString, QString> it(memberToGroupMap);
    while (it.hasNext())
    {
        it.next();
        groups[it.value()].addMember(members[it.key()]);
    }

    return groups.values().toVector().toStdVector();
}

SymmetricKeyInfoList DatabaseManager::allSymmetricKeys()
{
    QSqlQuery query(m_groupsDB);

    if (!query.exec("SELECT keyfiles.symkeyid, keyfiles.symkeyfilename, keyfiles.decryptonly, keyfiles.offset "
                    "FROM   keyfiles;"))
        throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, m_groupsDB.lastError().text() + ". Query: '" + query.lastQuery() + "'");

    SymmetricKeyInfoList list;

    while (query.next())
    {
        QString symkeyid          = query.value(0).toString();
        QString symkeyfile        = query.value(1).toString();
        bool decryptOnly          = query.value(2).toBool();
        unsigned long long offset = query.value(3).toULongLong();

        list.push_back(SymmetricKeyInfo(symkeyid.toStdString(), symkeyfile.toStdString(), decryptOnly, offset));
    }

    return list;
}

SymmetricKeyInfoList DatabaseManager::getSymmetricKeysAssignedToKeyID(const std::string& keyid)
{
    QSqlQuery query(m_groupsDB);

    if (!query.exec("SELECT groupmembers.keyid, groupmembers.symkeyid, keyfiles.symkeyid, keyfiles.symkeyfilename, keyfiles.decryptonly, keyfiles.offset "
                    "FROM   groupmembers, keyfiles "
                    "WHERE  groupmembers.keyid = '" + QString(keyid.c_str()) + "'" +
                    "AND    groupmembers.symkeyid = keyfiles.symkeyid;"))
        throw DatabaseManagerException(DatabaseManagerException::QUERY_ERROR, m_groupsDB.lastError().text() + ". Query: '" + query.lastQuery() + "'");

    SymmetricKeyInfoList list;

    while (query.next())
    {
        //QString keyid             = query.value(0).toString();
        QString symkeyid          = query.value(1).toString();
        QString symkeyfile        = query.value(3).toString();
        bool decryptOnly          = query.value(4).toBool();
        unsigned long long offset = query.value(5).toULongLong();

        list.push_back(SymmetricKeyInfo(symkeyid.toStdString(), symkeyfile.toStdString(), decryptOnly, offset));
    }

    return list;
}

}}
